from Bio import SeqIO
from Bio.Seq import Seq
import vcf
import argparse
import logging
import random
# 从vcf选择邻近的序列文件，生成fasta文件
# 读取参考基因组FASTA文件
#V3.0: DVR模式中，针对-链的情况，输出的DNA序列是5'-3'端的，而RNA模式的序列则是反向互补序列
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def read_fasta(fasta_file):
    logging.info("processing fasta...")
    sequences = {}
    with open(fasta_file, "r") as handle:
        for record in SeqIO.parse(handle, "fasta"):
            sequences[record.id] = record.seq
    logging.info("finish processing fasta")
    return sequences

# 提取变异位点周围的序列
def extract_sequences(ref_seq, variant_pos, flank_length=20):
    """
    提取变异位点周围的序列,并验证REF碱基是否正确。

    Parameters:
        ref_seq (str): 参考序列。
        variant_pos (int): 变异位点的位置。
        flank_length (int): 序列提取的侧翼长度。

    Returns:
        str: 提取出的序列。
    """
    logging.debug(f"Extracting sequence around position {variant_pos} with a flank length of {flank_length}.")
    start_pos = max(0, variant_pos - flank_length)
    end_pos = min(len(ref_seq), variant_pos + flank_length + 1)  # +1以包含变异位点本身
    extracted_seq = ref_seq[start_pos:end_pos]
    # 验证REF碱基
    ref_base = ref_seq[variant_pos]  # 调整索引以匹配提取序列的索引
    return extracted_seq, ref_base

# 读取VCF文件
def read_vcf(vcf_file):
    """
    从VCF文件中读取变异位点的信息

    Parameters:
        vcf_file (str): VCF文件的路径

    Returns:
        list: 包含变异位点信息的列表，每个元素是一个元组 (Chrom, Site, Ref)
    """
    logging.info("processing vcf...")
    variants = []
    with open(vcf_file, "r") as handle:
        vcf_reader = vcf.Reader(handle)
        for record in vcf_reader:
            Chrom = record.CHROM
            Site = record.POS - 1
            Ref = record.REF
            Strand = "+"
            variants.append((Chrom, Site, Ref, Strand))
    return variants

def read_chromsites(vcf_file):
    #这个函数的功能是从一个特定格式的文件（DVR结果文件）中读取变异位点的位置信息，
    #并将这些信息存储在一个列表中返回。这个函数假设文件的第一行可能是标题，而实际的数据行从第二行开始。
    chromsites = []
    with open(vcf_file, 'r') as file:
        # 读取第一行
        header = next(file).strip().split('\t')
        # 检查第一行的第一个元素是否为数字
        if header[0].isdigit():
            logging.info("The first line starts with a number, processing data...")
            # 因为第一行是数据行，所以直接处理它
            Chrom = header[0]
            Site = int(header[1]) - 1
            Ref = header[2]
            Strand = header[4]
            chromsites.append((Chrom, Site, Ref, Strand))
        logging.info("processing DVR data...")
        # 第一行被视为标题，现在开始从下一行读取数据
        for line in file:
            fields2 = line.strip().split('\t')
            Chrom = fields2[0]
            Site = int(fields2[1]) - 1
            Ref = fields2[2]
            Strand = fields2[4]
            chromsites.append((Chrom, Site, Ref, Strand))
    return chromsites
# 生成新的FASTA文件
def generate_random_fasta(output_file, ref_sequences, variants, num_variants=0, flank_length=20):
    """
    生成包含随机选择的变异位点序列的 FASTA 文件 并验证REF碱基。
    根据输入的绝对位置输出5'到3'的序列(参考),不管正负链
    Parameters:
        output_file (str): 输出的 FASTA 文件名。
        ref_sequences (dict): 参考序列字典。
        variants (list): 变异位点列表。由元组组成的列表,每个位点是一个元组,包含(Chrom, Site, Ref, Strand)
        num_variants (int): 随机选择的变异位点数量。
        flank_length (int): 序列提取的侧翼长度。

    Returns:
        None
    """
    logging.info(f"Preparing to generate FASTA file '{output_file}' with {num_variants if num_variants else 'all'} random variants and flank length of {flank_length}.")
    if num_variants:
        variants = random.sample(variants, num_variants)

    with open(output_file, "w") as handle:
        for chrom, pos, expected_ref, Strand in variants:
            ref_seq = ref_sequences[chrom]
            extracted_seq, actual_ref = extract_sequences(ref_seq, pos, flank_length)
            # 验证REF碱基是否匹配
            if expected_ref != actual_ref:
                logging.warning(f"Mismatch detected for variant at {chrom}:{pos+1}. Expected REF '{expected_ref}', found '{actual_ref}'.")
            handle.write(f">{chrom}_{pos+1}_{expected_ref}_{Strand}\n{extracted_seq}\n")
#只选定特定的序列，并且切换为RNA模式
def get_reverse_complement(sequence):
    #已经有了一个序列字符串，并且想要直接得到它的反向互补序列，可以将其封装进一个函数中：
    seq_obj = Seq(sequence)
    return str(seq_obj.reverse_complement())
def replace_t_with_u(input_fasta, output_fasta, flank_length):
    logging.info("RNA mode starting...")
    #RNA模式，将负链反转互补，且T变成U
    with open(input_fasta, "r") as infile, open(output_fasta, "w") as outfile:
        for line in infile:
            if line.startswith(">"):
                # 如果是标题行，直接写入输出文件
                sequence_id = line.strip()
                Strand = sequence_id.split('_')[3]
                outfile.write(sequence_id + "\n")
            else:
                if Strand == "-":
                    line = get_reverse_complement(line.strip())
                    modified_sequence = line.upper().replace("T", "U")
                    logging.debug("seq_-:"+modified_sequence)
                    outfile.write(modified_sequence + "\n")
                else:
                    modified_sequence = line.upper().replace("T", "U").strip()
                    logging.debug("seq_+:"+modified_sequence)
                    outfile.write(modified_sequence + "\n")

# 主函数
def main():
    parser = argparse.ArgumentParser(description='''
    这个程序用于将vcf文件或者DVR的结果文件位点的前后序列找到
    选择-v则是vcf模式,-d则是dvr结果模式
    如果有-U参数,设定为out文件的名称,则会输出一个将T转换为U的fasta文件
    默认位点前后长度20bp,可以通过-l自由设定
    2.0版本更新:增加了根据vcf文件的验证功能
    3.0版本: DVR模式中,针对-链的情况,输出的DNA序列是5'-3'端的m而RNA模式的序列则是反向互补序列(3'-5')
                                                ''')
    parser.add_argument("-f", "--fasta", help="Reference FASTA file")
    parser.add_argument("-v", "--vcf", help="VCF file")
    parser.add_argument("-d", "--dvr", help="results file for filtering DVR results")
    parser.add_argument("-o", "--output", help="Output FASTA file")
    parser.add_argument("-U", "--output_with_U", help="Output FASTA file with U")
    parser.add_argument("-l",'--length', type=int, default=20, help='Length parameter')
    parser.add_argument("-n", "--num_sequences", type=int, default=0, help="Number of random sequences to extract default: Do not random extract")

    args = parser.parse_args()
    ref_sequences = read_fasta(args.fasta)
    if args.dvr:
        variants = read_chromsites(args.dvr)
    elif args.vcf:
        variants = read_vcf(args.vcf)
    generate_random_fasta(args.output, ref_sequences, variants, args.num_sequences, args.length)

    if args.output_with_U:
        replace_t_with_u(args.output, args.output_with_U, args.length)
    logging.info("==============finish=================")


if __name__ == "__main__":
    main()
